use std::collections::{HashMap, HashSet};
use std::fs;

// 定义并查集结构
struct UnionFind {
    parent: Vec<usize>,
}

impl UnionFind {
    fn new(size: usize) -> Self {
        UnionFind {
            parent: (0..size).collect(),
        }
    }

    fn find(&mut self, x: usize) -> usize {
        if self.parent[x] != x {
            self.parent[x] = self.find(self.parent[x]);
        }
        self.parent[x]
    }

    fn union(&mut self, x: usize, y: usize) {
        let root_x = self.find(x);
        let root_y = self.find(y);
        if root_x != root_y {
            self.parent[root_y] = root_x;
        }
    }
}

struct CityConnections {
    connections: HashMap<usize, HashMap<String, Vec<String>>>,
}

fn load_city() -> CityConnections {
    fn parse_json(input: &str) -> HashMap<String, Vec<String>> {
        // let input = r#"
        // {
        //     "北京": ["房山"],
        //     "北京": ["大兴", "通州"],
        //     "北京": ["房山", "昌平"],
        //     "昌平": ["怀柔"]
        // }
        // "#;

        // 删除字符串中的 {} 和 \n
        let cleaned = input
            .replace("{", "")
            .replace("}", "")
            .replace(" ", "")
            .replace("\n", "");

        // 根据 ], 切割字符串
        let mut parts: Vec<String> = cleaned.split("],").map(|s| s.to_string()).collect();
        let mut parsed_maps: Vec<HashMap<String, Vec<String>>> = Vec::new();
        for s in &mut parts {
            if !s.ends_with(']') {
                s.push(']');
            }
            s.insert(0, '{');
            s.push('}');
            let parsed_map: HashMap<String, Vec<String>> = serde_json::from_str(s).unwrap();
            parsed_maps.push(parsed_map);
        }

        let mut city_parsed_maps: HashMap<String, Vec<String>> = HashMap::new();
        for part in parsed_maps {
            if let Some((key, value)) = part.iter().next() {
                if city_parsed_maps.contains_key(key) {
                    let mut set: HashSet<String> = value.clone().into_iter().collect();
                    set.extend(city_parsed_maps.get(key).unwrap().iter().cloned());
                    city_parsed_maps.insert(key.to_string(), set.into_iter().collect());
                } else {
                    city_parsed_maps.insert(key.clone(), value.clone());
                }
            }
        }
        return city_parsed_maps;
    }

    let mut cityconnections = CityConnections {
        connections: HashMap::new(),
    };
    // 读取 JSON 文件内容
    let json_str = fs::read_to_string("district.json").unwrap();

    let city_list = json_str.replace(" ", "").replace("\n", "");
    let mut city_parts: Vec<String> = city_list.split("},").map(|s| s.to_string()).collect();
    let mut city_connections: HashMap<usize, HashMap<String, Vec<String>>> = HashMap::new();
    for part in &mut city_parts {
        let part_split = part.split("\":{").collect::<Vec<&str>>();
        let id_str = part_split
            .get(0)
            .unwrap()
            .replace("{", "")
            .replace("\"", "")
            .parse::<usize>()
            .unwrap();
        let city_map = format!("{{{}}}", part_split.get(1).unwrap());
        city_connections.insert(id_str, parse_json(city_map.as_str()));
    }
    cityconnections.connections = city_connections;

    return cityconnections;
}

fn count_connected_components(graph: &HashMap<String, Vec<String>>) -> usize {
    // 收集所有城市名称
    let mut city_set: HashSet<String> = HashSet::new();
    for (city, neighbors) in graph {
        city_set.insert(city.clone());
        for neighbor in neighbors {
            city_set.insert(neighbor.clone());
        }
    }

    // 创建城市列表和映射
    let mut city_list: Vec<String> = Vec::new();
    for city in city_set {
        city_list.push(city);
    }
    let city_to_id: HashMap<String, usize> = city_list
        .iter()
        .enumerate()
        .map(|(id, city)| (city.clone(), id))
        .collect();

    // 初始化并查集
    let mut uf = UnionFind::new(city_list.len());

    // 进行 union 操作
    for (city, neighbors) in graph.clone() {
        let city_id = city_to_id[&city];
        for neighbor in neighbors {
            let neighbor_id = city_to_id[&neighbor];
            uf.union(city_id, neighbor_id);
        }
    }

    // 根据 find 结果分组
    let mut groups: HashMap<usize, Vec<String>> = HashMap::new();
    for city in &city_list {
        let root = uf.find(*city_to_id.get(city).unwrap());
        groups.entry(root).or_insert(Vec::new()).push(city.clone());
    }

    return groups.len();
}

pub fn count_provinces() -> String {
    let city_connections = load_city();

    let mut result: HashMap<usize, usize> = HashMap::new();
    for (id, city_map) in &city_connections.connections {
        let count = count_connected_components(city_map);
        result.insert(*id, count);
    }
    let mut result_list = Vec::new();

    for i in 1..=result.len() {
        result_list.push(result.get(&i).unwrap());
    }

    return result_list
        .iter()
        .map(|x| x.to_string())
        .collect::<Vec<String>>()
        .join(",");
}
